"""
Brief:
  模型评估相关，可以自己跑一下看看效果
Author:GL
"""


from model import ClassifierNet
from dataset import *


import torch
import cv2


def save_model(ckpt_path,output_path=None):
    '''
    Brief:
      从checkpoint中解析并保存模型
    Args:
      ckpt_path(str):checkpoint的模型路径，如./runs/checkpoints/xxx.pth.tar,里面还保存了学习率等数据
      output_path(str):模型的输出路径
    Return:None
    '''
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    model = ClassifierNet().to(device)
    model_dict = model.state_dict()

    checkpoint = torch.load(ckpt_path)
    #checkpoint是个压缩文件，不仅包括权重，因此需要单独写一个转换的函数
    model.load_state_dict(checkpoint['state_dict'])
    start_epoch = checkpoint['epoch']
    best_prec1 = checkpoint['best_prec1']
    print("epochs:{},prec1:{}".format(start_epoch,best_prec1))
    if output_path is None:
        torch.save(model,'./CNN_Classifier.pth',_use_new_zipfile_serialization=False)
    else :
        torch.save(model,output_path,_use_new_zipfile_serialization=False)

def test():
    model = torch.load('/home/lin/桌面/HITCRT/High_exo_CNN/CNN_Classifier.pth')
    model.eval()
    while True:
        path =input("img path:")
        if path =="q":
            break
        img = cv2.imread(path,0)
        #prepross
        img = cv2.resize(img,(52,52),interpolation=cv2.INTER_AREA)
        tensor = torch.from_numpy(img).permute(2, 0, 1).contiguous()
        tensor.unsqueeze(0)#append dim
        # end of preprocess
        pred = model.forward(tensor)
        conf,index = torch.max(pred,1)
        print("conf:{},class:{}".format(conf,pred))
    



if __name__ =='__main__':
   save_model('/home/hitcrt/桌面/HITCRT2021/numdetect/High_exo_CNN/runs/checkpoints/model_best.pth.tar')
   test()